package main

import (
	"crypto/md5"
	"fmt"
	"io"

	"math"
	"navidb-tool/config"
	"navidb-tool/logger"
	"navidb-tool/rocksdb"
	"sync"

	"github.com/panjf2000/ants/v2"
)

var (
	// servers = make(map[int]*rocksdb.RocksDBService)
	servers sync.Map

	mu sync.RWMutex
)

func main() {
	config.Init()
	logger.Init()
	createDB()
	serverNum := len(config.Conf.Rocks.DbPath)

	var swg sync.WaitGroup
	pool, _ := ants.NewPoolWithFunc(config.Conf.Rocks.Thread, writeTask)
	defer pool.Release()
	hds := config.Conf.Rocks.Hd
	for _, hd := range hds {
		length := len(fmt.Sprintf("%d", hd))
		start := hd * int(math.Pow10(11-length))
		end := (hd + 1) * int(math.Pow10(11-length))
		logger.Infof("写入数据：%v ~ %v \n", start, end)
		for i := start; i < end; i++ {
			swg.Add(1)
			if i%100000 == 0 {
				logger.Infof("已写入10000条数据：%v", i)
			}
			value := fmt.Sprintf("%d", i)
			pool.Invoke(&TaskData{
				key:       getMD5(value),
				value:     value,
				wg:        &swg,
				serverNum: serverNum,
			})
		}
	}
	swg.Wait()
}

type TaskData struct {
	key       []byte
	value     string
	wg        *sync.WaitGroup
	serverNum int
}

func writeTask(data any) {
	taskData := data.(*TaskData)
	index := getShard(taskData.key, taskData.serverNum)
	vv, ok := servers.Load(index)
	if !ok {
		logger.Error("未找到对应的服务器")
		taskData.wg.Done()
		return
	}
	ser := vv.(*rocksdb.RocksDBService)
	ser.Write(taskData.key, []byte(taskData.value))
	taskData.wg.Done()
}

func getMD5(key string) []byte {
	h := md5.New()
	io.WriteString(h, key)
	return h.Sum(nil)
}

func getShard(key []byte, num int) int {
	mu.RLock()
	defer mu.RUnlock()
	hash := md5.Sum(key)
	shardId := int(hash[0]) % num
	return shardId
}

func createDB() {
	for index, item := range config.Conf.Rocks.DbPath {
		logger.Infof(">>> Initializing %s RocksDB", item)
		service, err := rocksdb.NewRocksDBService(
			config.Conf.Rocks.NumShards,
			item,
		)
		if err != nil {
			logger.Errorf("Error initializing RocksDB : %v", err)
			logger.Panic(err)
		}
		logger.Infof("<<< RocksDB %s initialized", item)
		servers.Store(index, service)
	}
}
